#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2020/6/5 下午3:38
# @Author  : fugang_le
# @Software: PyCharm

# -*- coding: utf-8 -*-
"""
Created on Mon Jan  7 16:30:53 2019

@author: 003723
"""
import math
import sqlite3
import operator
import configparser
import pandas as pd

from app.libs.src.util import segment
from app.libs.log_util import logger
from app.libs.src.data_helper import stop_words


def is_number(s):
    try:
        float(s)
        return True
    except ValueError:
        return False


def clean_list(seg_list):
    cleaned_dict = {}
    n = 0
    for i in seg_list:
        i = i.strip().lower()
        if i != '' and not is_number(i) and i not in stop_words:
            n = n + 1
            if i in cleaned_dict:
                cleaned_dict[i] = cleaned_dict[i] + 1
            else:
                cleaned_dict[i] = 1
    return n, cleaned_dict


class IndexModule:
    def __init__(self, config_path):
        self.config_path = config_path
        self.config = configparser.ConfigParser()
        self.config.read(config_path)
        self.postings_lists = {}

    def write_postings_and_knowledge_to_db(self, db_path):
        conn = sqlite3.connect(db_path)
        c = conn.cursor()

        c.execute('''DROP TABLE IF EXISTS postings''')
        c.execute('''CREATE TABLE postings
                     (term TEXT PRIMARY KEY, df INTEGER, docs TEXT)''')

        c.execute('''DROP TABLE IF EXISTS knowledge''')
        c.execute('''CREATE TABLE knowledge
                     (id INTEGER PRIMARY KEY, question TEXT, answer TEXT)''')

        for key, value in self.postings_lists.items():
            doc_list = '\n'.join(map(str, value[1]))
            t = (key, value[0], doc_list)
            c.execute("INSERT INTO postings VALUES (?, ?, ?)", t)

        conn.commit()

        for i, question in self.files.items():
            answer = '标准问“' + question + '”的答案'
            t = (i, question, answer)
            c.execute("INSERT INTO knowledge VALUES (?, ?, ?)", t)

        conn.commit()
        conn.close()

    def construct_postings_lists(self, files, db_path):
        self.files = {k: v for k, v in enumerate(files)}

        AVG_L = 0

        for i, x in self.files.items():
            seg_list = segment(x)
            ld, cleaned_dict = clean_list(seg_list)
            AVG_L = AVG_L + ld

            for key, value in cleaned_dict.items():
                d = [i, value, ld]
                if key in self.postings_lists:
                    self.postings_lists[key][0] = self.postings_lists[key][0] + 1  # df++
                    self.postings_lists[key][1].append(d)
                else:
                    self.postings_lists[key] = [1, [d]]  # [df, [Doc]]
        AVG_L = AVG_L / len(files)
        self.config.set('DEFAULT', 'N', str(len(files)))
        self.config.set('DEFAULT', 'avg_l', str(AVG_L))
        with open(self.config_path, 'w') as configfile:
            self.config.write(configfile)
        self.write_postings_and_knowledge_to_db(db_path)


class SearchEngine:
    def __init__(self, config_path, db_path):
        self.config = configparser.ConfigParser()
        self.config.read(config_path)

        self.conn = sqlite3.connect(db_path, check_same_thread=False)
        self.K1 = float(self.config['DEFAULT']['k1'])
        self.B = float(self.config['DEFAULT']['b'])
        self.N = int(self.config['DEFAULT']['n'])
        self.AVG_L = float(self.config['DEFAULT']['avg_l'])

    def __del__(self):
        self.conn.close()

    def fetch_postings_db(self, term):
        c = self.conn.cursor()
        c.execute('SELECT * FROM postings WHERE term=?', (term,))
        return (c.fetchone())

    def fetch_knowledge_db(self, id):
        c = self.conn.cursor()
        c.execute('SELECT question FROM knowledge WHERE id=?', (id,))
        return (c.fetchone())

    def result_by_BM25(self, sentence, top_n=10):
        seg_list = segment(sentence)
        n, cleaned_dict = clean_list(seg_list)
        BM25_scores = {}
        for term in cleaned_dict.keys():
            r = self.fetch_postings_db(term)
            if r is None:
                continue
            df = r[1]
            w = math.log2((self.N - df + 0.5) / (df + 0.5))
            docs = r[2].split('\n')

            for doc in docs:
                docid, tf, ld = doc[1:-1].split(',')
                docid = int(docid)
                tf = int(tf)
                ld = int(ld)
                s = (self.K1 * tf * w) / (tf + self.K1 * (1 - self.B + self.B * ld / self.AVG_L))
                if docid in BM25_scores:
                    BM25_scores[docid] = BM25_scores[docid] + s
                else:
                    BM25_scores[docid] = s
        BM25_scores = sorted(BM25_scores.items(), key=operator.itemgetter(1))
        BM25_scores.reverse()
        if len(BM25_scores) == 0:
            return 0, []
        else:
            return 1, [(self.fetch_knowledge_db(x)[0], y) for x, y in BM25_scores][:top_n]